from tqdm import tqdm
import shelve
import time
import os
import json
import r2pipe
from bin2asm import *
import sys
sys.path.append(
    r"/home/cyw/projects/function_sim_project/siamese_Graphsage/asm2vec-pytorch-master/scripts")

# savePath=r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/radare2/"
savePath = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/radare2/"
# basePath=r"/home/cyw/projects/data/functionsim/allorigindata"
basePath = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/dataTrans"
# logPath=r"/home/cyw/projects/function_sim_project/all_logs/radare2_process_logs.txt"
logPath = r"/home/cyw/projects/function_sim_project/siamese_Graphsage/log.txt"


def write_logs(inf):
    with open(logPath, "a") as file:
        file.write(inf)


class radare2_wrapper():
    """
        实现radare2生成样本的图结构
        https://tower111.github.io/2021/03/13/radare%E8%8E%B7%E5%8F%96%E4%B8%80%E4%BA%9B%E4%BA%8C%E8%BF%9B%E5%88%B6%E6%96%87%E4%BB%B6%E4%BF%A1%E6%81%AF/
        学习radare2构建图结构
    """

    def __init__(self) -> None:
        pass

    def gene_sample_struct_inf(self, sampleName, fileName):
        """
            获得样本的图结构，和对应的汇编代码

            从保存的json中提取图结构
            保存的json文件中没有对应的callrefs.....难顶
            去看看radare2的官方文档，利用ax系列指令来处理
        """
        # Use -2 flag for non-interactive mode
        r2 = r2pipe.open(sampleName, flags=['-2'])
        # a越多分析的东西越多，asm2vec-pytorch用四个a，但是好像最多就三个
        # 分析二进制文件,
        r2.cmd('aaa')
        # 获得函数列表
        functions = r2.cmdj('aflj')
        # 保存获得的json文件,好像也不需要保存
        # https://www.json.cn/,这个网站展示比较清晰  json tools ctrl+alt+m
        saveName = savePath+fileName+".json"
        os.makedirs(os.path.dirname(saveName), exist_ok=True)
        with open(saveName, "w", encoding="utf-8") as file:
            json.dump(functions, file, ensure_ascii=False)

        res = {}
        # 函数间的调用
        calls = {}
        name_to_id = {}
        cnt = 0
        for function in functions:
            name = function["name"]
            calls[name] = []
            name_to_id[name] = cnt
            cnt += 1
        res["name_to_id"] = name_to_id
        res["size"] = cnt
        for function in functions:
            name = function["name"]
            address = function["offset"]
            # axf 该函数引用的所有地方(什么结果都没有，奇怪)  axt引用了对应地址的地方
            xrefs = r2.cmdj(f"axtj {address}")
            for xref in xrefs:
                if xref["type"] == "CALL":  # GPT乱搞，还得靠自己
                    # 加@表示地址，否则全为空
                    called_function_name = r2.cmd(
                        f"afn @{xref['from']}").strip()
                    # 特例处理
                    if called_function_name == "":
                        continue
                    # 这里反着写，逻辑才对  calls[name].append(called_function_name)
                    # 新问题修复,2个样本会出现keyError的情况,这里做个特例处理
                    if called_function_name in calls:
                        calls[called_function_name].append(name)
        res["calls"] = calls
        # 生成邻接矩阵
        matrix = [[0] * cnt for i in range(cnt)]
        for name in calls.keys():
            ind = name_to_id[name]
            for subName in calls[name]:
                if subName == "":
                    continue
                indy = name_to_id[subName]
                matrix[ind][indy] += 1
        res["adj"] = matrix
        discode = {}
        # 还需要获得每个函数的汇编代码
        for function in functions:
            r2.cmd(f's {function["offset"]}')
            asm = fn2asm(r2.cmdj('pdfj'), 0)
            if asm:
                # 不太确定格式是否一致！！！
                asm = f''' .name {function["name"]}\n.offset {function["offset"]:016x}\n''' + asm
            discode[function["name"]] = asm
        res["discode"] = discode
        return res

    def setter(self, data, name):
        """
            保存字典
            key如下：
                size:样本中函数的个数
                discode:函数名及其对应的反汇编代码
                adj:邻接矩阵
                calls：函数间调用关系，都是函数名
                name_to_id:函数名与下标的映射
        """
        with shelve.open(savePath+name+"_origin") as file:
            file["data"] = data
            file.close()

    def getter(self, name):
        with shelve.open(savePath+name) as file:
            data = file["data"]
            file.close()
        return data


if __name__ == "__main__":
    # 初步写完，吧整个流程跑通，确认没问题后，在进行样本的预处理
    # nohup python /home/cyw/projects/function_sim_project/siamese_Graphsage/gene_graph_struct_by_radare2.py
    radare2Tools = radare2_wrapper()
    # 获得已经处理好的样本名
    finishedName = {}
    cnt = 0
    g = os.walk(savePath)
    for path, dir_list, file_list in g:
        for fileName in tqdm(file_list):
            name, hz = fileName.split(".")
            if hz == "dir":
                finishedName[name.split("_")[0]] = True
                cnt += 1
    print("已经处理了{}个样本\n".format(cnt))
    cnt = 0
    error = 0
    errorName = []
    print("radara2 开始反汇编\n")
    g = os.walk(basePath)
    for path, dir_list, file_list in g:
        for fileName in tqdm(file_list):
            if fileName not in finishedName:
                cnt += 1
                filePath = os.path.join(path, fileName)
                a = time.time()
                write_logs("第{}个样本：路径{}\n".format(cnt, filePath))

                try:
                    res = radare2Tools.gene_sample_struct_inf(
                        filePath, fileName)
                    radare2Tools.setter(res, fileName)
                except Exception as e:
                    print(e)
                    error += 1
                    errorName.append(fileName)
                    write_logs("    发生异常{} s\n".format(e))
                finally:
                    b = time.time()
                    write_logs("    样本处理完成,耗时:{} s\n".format(b-a))
                    print("第{}个样本\n".format(cnt))
    print("任务完成~~\n")
    print("异常样本{}个\n".format(error))
    print(errorName)

    # res=radare2Tools.gene_sample_struct_inf('/home/cyw/projects/data/functionsim/allorigindata/742a22390bc72ec38eb18c95446972e8',"742a22390bc72ec38eb18c95446972e8")
    # radare2Tools.setter(res,"742a22390bc72ec38eb18c95446972e8")
    # print("--")
